from padding import Padding
import cv2
import numpy as np
from filter import Filter
from noise import Noise


def InvFilter(g, psf):

    G = np.fft.fft2(g)

    H = np.fft.fft2(psf)

    F_fit = G / H

    f_fit = np.fft.ifft2(F_fit)

    f_fit = np.fft.fftshift(f_fit)

    f_fit = np.abs(f_fit)

    f_fit = ((f_fit - f_fit.min()) / (f_fit.max() - f_fit.min()) * 255).astype('uint8')

    return f_fit


def WinnerFilter(g, psf, K):

    G = np.fft.fft2(g)

    H = np.fft.fft2(psf)

    F_fit = G * np.conj(H) /(np.abs(H)**2 + K)

    f_fit = np.fft.ifft2(F_fit)

    f_fit = np.fft.fftshift(f_fit)

    f_fit = np.abs(f_fit)

    f_fit = ((f_fit - f_fit.min()) / (f_fit.max() - f_fit.min()) * 255).astype('uint8')

    return f_fit


def MAP(g, psf, K):

    G = np.fft.fft2(g)

    H = np.fft.fft2(psf)

    h, w = g.shape

    C = np.arange(0, h).reshape((-1, 1)) ** 2 + np.arange(0, w) ** 2

    F_fit = G * np.conj(H) /(np.abs(H)**2 + K * C)

    f_fit = np.fft.ifft2(F_fit)

    f_fit = np.fft.fftshift(f_fit)

    f_fit = np.abs(f_fit)

    f_fit = ((f_fit - f_fit.min()) / (f_fit.max() - f_fit.min()) * 255).astype('uint8')

    return f_fit


def MLE(g, h, n):

    f_fit = g

    for _ in range(n):
        f_fit = f_fit * cv2.filter2D(g / cv2.filter2D(f_fit, -1, h), -1, h)

    f_fit = ((f_fit - f_fit.min()) / (f_fit.max() - f_fit.min()) * 255).astype('uint8')

    return f_fit


if __name__ == '__main__':
    SRC = cv2.imread('peppers-bw.bmp', cv2.IMREAD_GRAYSCALE)

    def loss(y, label):
        return 10 * np.log10(255 ** 2 / np.mean((label - y) ** 2 + 1e-6))

    def show(winname, img):
        cv2.imshow(winname, img)
        cv2.resizeWindow(winname, 512, 512)
        cv2.waitKey(1)
        print(winname, str(loss(SRC, img)))

    show('SRC', SRC)

    # 退化模型
    degenerate = Filter(Padding('REAPEAT'), 'GAUSS')
    # 噪声模型
    noise = Noise()
    # 参数
    dsize, dsigma, sigma = 19, 2, 5
    # 退化图像
    g = noise(degenerate(SRC, size=dsize, sigma=dsigma), sigma=sigma)
    
    G = np.fft.fft2(g)
    # 模板核
    kernel = degenerate.get_kernel()
    # 点扩散函数
    psf = np.zeros_like(SRC, dtype='float64')

    h, w = SRC.shape

    psf[h // 2 - dsize // 2 : h // 2 + dsize // 2 + 1, w // 2 - dsize // 2 : w // 2 + dsize // 2 + 1] = kernel

    show('GUASS', g)

    show('INV', InvFilter(g, psf))

    show('Winner', WinnerFilter(g, psf, K=0.03))

    show('MAP', MAP(g, psf, K=1e-6))

    show('MLE', MLE(g, kernel, 10))

    cv2.waitKey(0)


"""
参考：
https://blog.csdn.net/bingbingxie1/article/details/79398601
https://blog.51cto.com/nailaoer/2896637
https://blog.csdn.net/wsp_1138886114/article/details/97683219
"""